"""DeepSeek模型适配器实现。"""

from typing import AsyncGenerator, Dict, Any, List
import aiohttp
import json

from .base import BaseModelAdapter, ModelConfig, ModelResponse


class DeepSeekAdapter(BaseModelAdapter):
    """DeepSeek大模型适配器，专注于代码生成。"""
    
    # 定价（每千tokens，美元）
    PRICING = {
        "deepseek-coder": {"prompt": 0.0005, "completion": 0.0015},
        "deepseek-chat": {"prompt": 0.0005, "completion": 0.0015}
    }
    
    def _initialize_client(self) -> None:
        """初始化DeepSeek客户端。"""
        self.base_url = self.config.api_base or "https://api.deepseek.com/v1"
    
    async def generate(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> ModelResponse:
        """生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Returns:
            模型响应
        """
        url = f"{self.base_url}/chat/completions"
        
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }
        
        # DeepSeek特有参数
        code_mode = kwargs.pop("code_mode", False)
        
        data = {
            "model": self.config.model_name,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p,
            "frequency_penalty": self.config.frequency_penalty,
            "presence_penalty": self.config.presence_penalty,
            "stream": False
        }
        
        # 代码模式优化
        if code_mode:
            data["temperature"] = min(0.3, self.config.temperature)  # 降低温度以提高代码准确性
            data["stop"] = kwargs.get("stop", ["\n\n\n", "```\n\n"])  # 代码块结束标记
        
        data.update(kwargs)
        data.update(self.config.extra_params)
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers=headers,
                json=data,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                result = await response.json()
                
                if "error" in result:
                    raise Exception(f"DeepSeek API错误: {result['error'].get('message', 'Unknown error')}")
                
                choice = result["choices"][0]
                usage = result.get("usage", {})
                
                # 代码后处理
                content = choice["message"]["content"]
                if code_mode:
                    content = self._format_code(content)
                
                return ModelResponse(
                    content=content,
                    model=result.get("model", self.config.model_name),
                    usage={
                        "prompt_tokens": usage.get("prompt_tokens", 0),
                        "completion_tokens": usage.get("completion_tokens", 0),
                        "total_tokens": usage.get("total_tokens", 0)
                    },
                    finish_reason=choice.get("finish_reason", "stop"),
                    metadata={
                        "id": result.get("id", ""),
                        "code_mode": code_mode
                    }
                )
    
    async def stream(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """流式生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Yields:
            内容块
        """
        url = f"{self.base_url}/chat/completions"
        
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }
        
        code_mode = kwargs.pop("code_mode", False)
        
        data = {
            "model": self.config.model_name,
            "messages": messages,
            "temperature": min(0.3, self.config.temperature) if code_mode else self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p,
            "stream": True
        }
        
        data.update(kwargs)
        data.update(self.config.extra_params)
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers=headers,
                json=data,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                async for line in response.content:
                    if line:
                        line_str = line.decode('utf-8').strip()
                        if line_str.startswith("data: "):
                            data_str = line_str[6:]
                            if data_str == "[DONE]":
                                break
                            
                            try:
                                chunk = json.loads(data_str)
                                if "choices" in chunk and chunk["choices"]:
                                    delta = chunk["choices"][0].get("delta", {})
                                    content = delta.get("content", "")
                                    if content:
                                        yield content
                            except json.JSONDecodeError:
                                continue
    
    def _format_code(self, content: str) -> str:
        """格式化代码输出。
        
        Args:
            content: 原始内容
            
        Returns:
            格式化后的内容
        """
        # 简单的代码格式化，实际可以更复杂
        lines = content.split('\n')
        formatted_lines = []
        in_code_block = False
        
        for line in lines:
            if line.strip().startswith('```'):
                in_code_block = not in_code_block
                formatted_lines.append(line)
            elif in_code_block:
                # 保持代码块内的格式
                formatted_lines.append(line)
            else:
                # 非代码块内容
                formatted_lines.append(line)
        
        return '\n'.join(formatted_lines)
    
    def count_tokens(self, text: str) -> int:
        """计算文本的token数量。
        
        DeepSeek使用类似GPT的tokenizer。
        
        Args:
            text: 输入文本
            
        Returns:
            Token数量
        """
        # 简单估算：平均4个字符1个token
        return len(text) // 4
    
    def estimate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """估算使用成本。
        
        Args:
            prompt_tokens: 提示词token数
            completion_tokens: 生成token数
            
        Returns:
            预估成本（美元）
        """
        model_key = "deepseek-coder" if "coder" in self.config.model_name else "deepseek-chat"
        pricing = self.PRICING[model_key]
        
        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        
        return prompt_cost + completion_cost
    
    async def generate_code(
        self,
        prompt: str,
        language: str = "python",
        **kwargs
    ) -> ModelResponse:
        """专门的代码生成方法。
        
        Args:
            prompt: 代码生成提示
            language: 编程语言
            **kwargs: 额外参数
            
        Returns:
            包含代码的响应
        """
        # 构建专门的代码生成提示
        system_message = f"You are an expert {language} programmer. Generate clean, efficient, and well-commented code."
        
        messages = [
            {"role": "system", "content": system_message},
            {"role": "user", "content": prompt}
        ]
        
        return await self.generate(messages, code_mode=True, **kwargs)